#!/usr/bin/env python
"""

   align_gaps.py

        handles sequences given a gap string
        specified in the following format "gap_pos:gap_len|gap_pos:gap_len|gap_pos:gap_len"
 


   Copyright (C) 2007 Leo Goodstadt

   This program is free software; you can redistribute it and/or
   modify it under the terms of the GNU General Public License
   as published by the Free Software Foundation; version 2
   of the License

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
"""


import sys

# add self to search path for testing
if __name__ == '__main__':
    sys.path.append("../python_modules")

from general_util import die_error
from iterator_helper import count_while


#_________________________________________________________________________________________

#   gapstr_to_gaps

#       input:

#           gap_str
#           e.g. 0:63|119:6|263:1

#       returns:

#            [[beg_pos:length][beg_pos:length]]

#   gaps_to_gapstr

#       reverses above

#_________________________________________________________________________________________
def gapstr_to_gaps(gap_str):
    if not len(gap_str):
        return list()
    return [map(int, gap.split(':')) for gap in gap_str.split("|")]

def gaps_to_gapstr(gaps):
    return "|".join ([":".join(map(str, gap)) for gap in gaps])
        

#_________________________________________________________________________________________

#   expand_cdna_with_gaps

#       input:

#           seq
#           aligning character
#           gaps: list of position/length pairs in string format 0:63|119:6|263:1

#       returns:

#           $expanded sequences
#_________________________________________________________________________________________
def expand_cdna_with_gaps(seq, gapchar = '_', gaps = []):
    gapchar = (gapchar[0:1]) * 3
    for gap in reversed(gaps):
        # insert a gap of gap[1] * 3 size at gap[0] * 3 position
        seq = seq[: gap[0] * 3] + gapchar * gap[1] + seq[gap[0] * 3:]
    return seq
        
#_________________________________________________________________________________________

#   expand_prot_with_gaps

#       input:

#           $seq
#           $aligning character
#           @list of position/length pairs

#       returns:

#           $expanded sequences
#_________________________________________________________________________________________
def expand_prot_with_gaps(seq, gapchar = '_', gaps = []):
    gapchar = gapchar[0:1]
    for gap in reversed(gaps):
        # insert a gap of gap[1] size at gap[0] position
        seq = seq[:gap[0]] + gapchar * gap[1] + seq[gap[0]:]
    return seq



#_________________________________________________________________________________________

#   helper functions for gap strings

#_________________________________________________________________________________________
def expand_cdna_with_gapstr(seq, gapchar = '_', gapstr = ''):
    return expand_cdna_with_gaps(seq, gapchar, gapstr_to_gaps(gapstr))
def expand_prot_with_gapstr(seq, gapchar = '_', gapstr = ''):
    return expand_prot_with_gaps(seq, gapchar, gapstr_to_gaps(gapstr))
    



#_________________________________________________________________________________________

#   map_pos_to_aligned

#       input:

#           position
#           gaps

#       returns:

#           adjusted position given the aligning marks
#_________________________________________________________________________________________
def map_pos_to_aligned(pos, gaps):
    resulting_pos = pos
    for gap in gaps:
        if pos >= gap[0]:
            resulting_pos += gap[1]
        else:
            break
    #    align2 = [[0,4],[11,3],[14,3]]
    return resulting_pos

#_________________________________________________________________________________________

#   map_pos_from_aligned

#       input:

#           position
#           gaps

#       returns:

#           adjusted position given the aligning marks
#_________________________________________________________________________________________
def map_pos_from_aligned(pos, gaps):
    for gap in gaps:
        if gap[0] >= pos:
            if gap[0] == pos:
                return max(pos - 1, 0)
            break

        pos -= gap[1]

        # in gap region
        if gap[0] > pos:
            pos = gap[0] - 1
            break

    return max(pos,  0)

#_________________________________________________________________________________________

#   map_pos_via_aligned


#       map from one alignment to another


#       input:

#           position
#           gaps of alignment position/length pairs for my sequence
#           gaps of alignment position/length pairs for other sequence

#       returns:

#           adjusted position given the aligning marks
#_________________________________________________________________________________________
def map_pos_via_aligned(pos, gaps1, gaps2):
    return map_pos_from_aligned(map_pos_to_aligned(pos, gaps1),  gaps2);













#_________________________________________________________________________________________

#   aligned_seq_to_gaps

#       input:

#           $ aligned protein sequence with gaps

#       returns:

#           positions of aligning marks
#_________________________________________________________________________________________
def aligned_seq_to_gaps(seq, gapchar = '-'):
    gaps = list()
    
    # save the number and position of the aligning characters for each sequence
    pos = seq.find(gapchar, 0)
    cummulative  = 0
    while pos <> -1:
        gap_len = count_while(x ==gapchar for x in iter(seq[pos:]))
        gaps.append([pos - cummulative, gap_len])
        cummulative += gap_len
        pos = seq.find(gapchar, pos + gap_len)

    return gaps



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
#
#
#   Test
#

def unit_test ():
    gapstr = "0:1199|172:276|501:68"
    cdna = "ATGACAGTGGAGCAGAAATTTGGCCTGTTTTCTGCTGAGATAAAGGAAGCAGACCCCCTGGCTGCCTCGGAAGCAAGTCAACCCAAACCCTGTCCCCCCGAAGTGACCCCCCACTACATCTGGATCGACTTCCTGGTGCAGCGGTTTGAGATCGCCAAGTACTGCAGCTCTGACCAAGTGGAGATCTTCTCCAGCCTGCTGCAGCGCTCCATGTCCCTGAACATCGGCAGGGCCAAGGGGAGCATGAACCGGCACGTGGCGGCCATCGGGCCCCGCTTCAAGCTGCTGACCCTGGGGCTGTCCCTCCTGCATGCCGATGTGGTTCCAAATGCAACCATCCGCAATGTGCTTCGCGAGAAGATCTACTCCACTGCCTTTGACTACTTCAGCTGTCCCCCAAAGTTTCCTACTCAAGGAGAGAAGCGGCTGCGTGAAGACATAAGCATCATGATTAAATTTTGGACCGCCATGTTCTCAGATAAGAAGTACCTGACCGCCAGCCAGCTTGTTCCCCCAGCTGACATCGGCGACCTCCTGGAGCAGTTAGTAGAGGAGAACACAGGCTCCTTGTCGGGCCCAGCGAAGGACTTTTACCAGCGGNNNTTTGATTTCTTTAACAAGATCACCAACGTGTCGGCTGTCATCAAGCCCTACCCTAAAGGCGACGAGAGAAAGAAGGCTTGTCTGTCGGCCCTGTCTGAAGTGACGGTGCAGCCAGGCTGCTCCCTGCCCAGCAACCCCGAAGCCATTGTGCTGGACGTCGACTACAAGTCTGGGACCCCGATGCAGAGTGCTGCAAAAGCCCCATATCTGGCCAAGTTCAAGGTGAAGCGATGTGGAGTTAGTGAACTTGAAAAAGAAGGTCTGCGGTGCCGCTCAGACTCTGAGGATGAGTGCAGCACGCAGGAGGCCGACGGCCAGAAGATCTCCTGGCAGGCAGCCATCTTCAAACTGGGAGACGACTGCCGGCAGGACATGCTGGCCCTGCAGATCATCGACCTCTTCAAGAACATCTTCCAGCTGGTCGGCCTGGACCTCTTTGTTTTTCCCTACCGCGTGGTGGCCACTGCCCCTGGGTGCGGGGTGATCGAGTGCATCCCCGACTGCACCTCCCGGGACCAGCTGGGCCGCCAGACAGACTTCGGCATGTACGACTACTTCACACGCCAGTACGGGGATGAGTCCACTCTGGCCTTCCAGCAGGCCCGCTACAACTTCATCCGAAGCATGGCCGCCTACAGCCTCCTGCTGTTCCTGCTGCAGATCAAGGACAGACACAACGGCAACATTATGCTGGACAAGAAGGGCCATATCATCCACATCGACTTTGGCTTCATGTTTGAAAGCTCGCCGGGCGGCAATCTCGGCTGGGAACCCGACATCAAGCTGACGGATGAGATGGTGATGATCATGGGGGGCAAGATGGAGGCCACACCCTTCAAGTGGTTCATGGAGATGTGTGTCCAGGCTACCTGGCTGTGCGGCACAGGTTTAGCCCCAACA"
    print expand_cdna_with_gapstr(cdna,  '_',  gapstr)
    
def unit_test1():
    PRINT_LEN = 25
    align1 = [[11,3]]
    align2 = [[0,4],[11,3],[14,3]]
    seq1 = 'A' * 21
    seq2 = 'B' * 14
    
    expand_prot_with_gapstr(seq1,  '-',  '')
    expand_prot_with_gapstr(seq2,  '-',  '')
    
    
    #
    #   Sanity checks
    #
    for x in range(50):
        print "1:", x, map_pos_to_aligned(x,  align1), map_pos_from_aligned(map_pos_to_aligned(x,  align1), align1)
        print "2:", x, map_pos_to_aligned(x,  align2), map_pos_from_aligned(map_pos_to_aligned(x,  align2), align2)
        assert map_pos_from_aligned(map_pos_to_aligned(x,  align1), align1) == x
        assert map_pos_from_aligned(map_pos_to_aligned(x,  align2), align2) == x

    print expand_prot_with_gaps(seq1, '-', align1)
    print expand_prot_with_gaps(seq2, '-', align2)
    print aligned_seq_to_gaps(expand_prot_with_gaps(seq1, '-', align1))
    print aligned_seq_to_gaps(expand_prot_with_gaps(seq2, '-', align2))
    print gaps_to_gapstr(aligned_seq_to_gaps(expand_prot_with_gaps(seq1, '-', align1)))
    print gaps_to_gapstr(aligned_seq_to_gaps(expand_prot_with_gaps(seq2, '-', align2)))
    assert(gaps_to_gapstr(aligned_seq_to_gaps(expand_prot_with_gaps(seq1, '-', align1))) == 
                gaps_to_gapstr(align1))
    assert(gaps_to_gapstr(aligned_seq_to_gaps(expand_prot_with_gaps(seq2, '-', align2))) == 
                gaps_to_gapstr(align2))
    
    
    #
    #   print conversion maps
    #
    print "\nseq1\t-> seq2\n----\t   ----\n"
    for x in range(26):
        print  str(x) + "\t-> " + str(map_pos_via_aligned(x, align1, align2))
    
    
    expanded1 = expand_prot_with_gaps(seq1, '-', align1)
    expanded2 = expand_prot_with_gaps(seq2, '-', align2)
    # convert to numbers
    def number_nongap_positions(seq):
        counter = 0
        numstr = ''
        for x in seq:
            if x == '-':
                numstr += ' '
            else:
                numstr += str((counter % 10))
                counter += 1
        return numstr
    
    print "\nSequences:"
    print number_nongap_positions(expanded1)[0:PRINT_LEN]
    print expanded1[0:PRINT_LEN]
    print ""
    for x in range(PRINT_LEN):
        sys.stdout.write (str(x % 10))
    print "\n"
    print expanded2[0:PRINT_LEN]
    print number_nongap_positions(expanded2)[0:PRINT_LEN]

        
                
if __name__ == '__main__':
    unit_test()
    



